# -*- coding: utf-8 -*-
"""
圆柱绕流势流场可视化程序

功能特性：
- 实时交互式参数调整（圆柱半径、来流速度）
- 基于势流理论的流场数值计算
- Numba加速核心算法
- PyQt5构建图形界面
- Matplotlib实现专业可视化
- 输入验证与物理验证机制
"""

# ================== 基础库导入 ==================
import numpy as np         # 数值计算核心库，用于矩阵运算和网格生成
import numba as nb         # 即时编译器，用于加速数值计算核心算法

# ================== GUI框架相关库 ==================
from PyQt5.QtWidgets import (
    QApplication,   # 应用管理类（管理事件循环等）
    QMainWindow,    # 主窗口基类
    QVBoxLayout,    # 垂直布局管理器
    QWidget,        # 基础控件容器
    QSlider,        # 滑动条控件
    QLineEdit,      # 文本输入框
    QLabel          # 文本标签
)
from PyQt5.QtCore import Qt        # Qt核心功能（包含方向常量等）
from PyQt5.QtGui import QValidator # 输入验证基类

# ================== 可视化库导入 ==================
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas  # Qt5后端画布
from matplotlib.figure import Figure  # 图表容器类
import matplotlib.pyplot as plt       # 绘图接口，提供流线图、等势线等可视化功能


# ================== 物理参数和计算核心 ==================
nu = 1.5e-5  # 空气的运动粘性系数 (m²/s)，标准大气条件15℃时的典型值

@nb.njit(parallel=True)
def compute_velocity(X, Y, U, a):
    """
    使用Numba加速的势流速度场计算（圆柱绕流解析解）
    
    参数:
        X, Y (ndarray): 网格坐标矩阵，形状相同的二维数组
        U (float): 来流速度，单位m/s
        a (float): 圆柱半径，单位m
        
    返回:
        Vx, Vy (ndarray): x和y方向的速度分量矩阵，与输入网格同形
        
    算法:
        基于势流理论解析解推导速度场：
        1. 复势函数：F(z) = U(z + a²/z)
        2. 复速度：V = dF/dz = U(1 - a²/z²)
        3. 分解实部虚部得到各点速度分量
    """
    Vx = np.zeros_like(X)
    Vy = np.zeros_like(Y)
    a_sq = a**2  # 圆柱半径平方提前计算
    
    # 并行遍历网格点（Numba自动优化循环）
    for i in nb.prange(X.shape[0]):
        for j in nb.prange(X.shape[1]):
            x = X[i, j]
            y = Y[i, j]
            r_sq = x**2 + y**2  # 计算到原点的距离平方
            
            # 排除原点附近和圆柱内部的点（避免奇异值）
            if r_sq < 1e-6 or r_sq <= a_sq:
                continue  # 保持速度为零
                
            z_sq_inv = 1 / r_sq  # 1/(x²+y²)
            # 分解复速度的实部和虚部
            real_part = 1 - a_sq*(x**2 - y**2)*z_sq_inv**2
            imag_part = 2*a_sq*x*y*z_sq_inv**2
            
            Vx[i, j] = U * real_part  # x方向速度分量
            Vy[i, j] = -U * imag_part  # y方向速度分量（负号保证速度方向正确）

    return Vx, Vy

def calculate_stream_function(X, Y, U, a):
    """
    计算势流流函数（圆柱绕流解析解）
    
    参数:
        X, Y (ndarray): 网格坐标矩阵，形状相同的二维数组
        U (float): 来流速度，单位m/s
        a (float): 圆柱半径，单位m
        
    返回:
        psi (ndarray): 流函数值矩阵，与输入网格同形
        
    注意:
        当r=0时添加微小值(1e-10)避免除以零错误，保持计算稳定性
        流函数公式：ψ = U*y*(1 - a²/(x²+y²))
    """
    r_sq = X**2 + Y**2
    safe_r_sq = np.where(r_sq != 0, r_sq, 1e-10)  # 处理原点附近的奇异点
    return U * Y * (1 - a**2 / safe_r_sq)  # 势流理论标准解


# ================== 输入验证器 ==================
class FloatValidator(QValidator):
    """
    自定义浮点数验证器（限制输入范围1.0~10.0）
    
    特性:
        - 实时验证输入有效性（绿色/红色边框反馈）
        - 自动修正超出范围的数值到边界值
        - 拒绝非数字输入
        - 支持中间状态（允许空输入）
    """
    def __init__(self, min_val, max_val, parent=None):
        super().__init__(parent)
        self.min_val = min_val  # 最小值边界
        self.max_val = max_val  # 最大值边界

    def validate(self, input_str, pos):
        """实时验证输入内容"""
        if not input_str:
            # 允许空输入作为中间状态（方便用户编辑）
            return (QValidator.Intermediate, input_str, pos)
        
        try:
            value = float(input_str)
            if self.min_val <= value <= self.max_val:
                return (QValidator.Acceptable, input_str, pos)
            return (QValidator.Invalid, input_str, pos)
        except ValueError:
            return (QValidator.Invalid, input_str, pos)

    def fixup(self, input_str):
        """自动修正非法输入为最近边界值"""
        try:
            value = float(input_str)
            value = np.clip(value, self.min_val, self.max_val)
            return f"{value:.1f}"  # 格式化为1位小数
        except:
            return str(self.min_val)  # 默认返回最小值


# ================== 主界面类 ==================
class MainWindow(QMainWindow):
    """
    主窗口类，负责GUI布局和可视化更新
    
    功能特性:
        - 实时交互式参数调整（滑动条+输入框）
        - 流场动态可视化（流线+等势线）
        - 雷诺数自动计算与显示
        - 流线闭合性物理验证
        - 界面样式美化
    """
    def __init__(self):
        super().__init__()
        # 初始化物理参数
        self.a = 1.0  # 圆柱半径 (m)，对应滑动条中间位置
        self.U = 5.0  # 来流速度 (m/s)，默认值
        
        self.init_ui()      # 初始化界面组件
        self.update_plot()  # 首次绘制

    def init_ui(self):
        """初始化用户界面组件和布局"""
        self.setWindowTitle("流场可视化系统 v2.0")
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # 圆柱半径控制滑动条（范围0.5~2.0m，精度0.1m）
        self.slider_a = QSlider(Qt.Horizontal)
        self.slider_a.setRange(5, 20)  # 内部使用整数避免浮点误差，实际值需除以10
        self.slider_a.setValue(10)     # 初始值对应1.0m
        self.slider_a.valueChanged.connect(self.update_a)
        layout.addWidget(QLabel("圆柱半径 (m):"))
        layout.addWidget(self.slider_a)

        # 流速输入框（带范围验证）
        self.lineEdit_U = QLineEdit(str(self.U))
        self.lineEdit_U.setValidator(FloatValidator(1.0, 10.0, self))
        self.lineEdit_U.editingFinished.connect(self.update_U)
        layout.addWidget(QLabel("流速 U (m/s):"))
        layout.addWidget(self.lineEdit_U)

        # 状态显示栏（雷诺数和验证结果）
        self.status_label = QLabel()
        layout.addWidget(self.status_label)

        # Matplotlib绘图区域设置
        self.figure = Figure(figsize=(8, 6), facecolor='black')  # 黑色背景
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)

        # 界面样式美化
        self.setStyleSheet("""
            QMainWindow { background: #333; }
            QLabel { color: #EEE; font: 10pt Arial; }
            QSlider::groove:horizontal {
                background: #555;
                height: 8px;
                border-radius: 4px;
            }
            QSlider::handle:horizontal {
                background: #FF8C00;
                width: 16px;
                margin: -4px 0;
                border-radius: 8px;
            }
            QLineEdit {
                border: 2px solid #888;
                border-radius: 4px;
                padding: 2px;
                background: #444;
                color: #FFF;
            }
            QLineEdit:valid { border-color: #0F0; }
            QLineEdit:invalid { border-color: #F00; }
        """)
        self.resize(1000, 800)  # 初始窗口尺寸

    def update_a(self):
        """处理圆柱半径变化事件（滑动条回调）"""
        self.a = self.slider_a.value() / 10.0  # 将滑动条整数转换为浮点值
        self.update_plot()

    def update_U(self):
        """处理流速输入变化事件（文本框回调）"""
        try:
            self.U = float(self.lineEdit_U.text())
            self.U = np.clip(self.U, 1.0, 10.0)  # 强制限制在有效范围
            self.lineEdit_U.setText(f"{self.U:.1f}")  # 格式化显示
        except ValueError:
            pass  # 忽略无效输入
        self.update_plot()

    def update_plot(self):
        """更新流场可视化图形（核心绘图方法）"""
        # 生成计算网格（范围：±5m，分辨率400×400）
        X, Y = np.meshgrid(np.linspace(-5, 5, 400),
                          np.linspace(-5, 5, 400))
        
        # 核心计算（Numba加速）
        Vx, Vy = compute_velocity(X, Y, self.U, self.a)
        psi = calculate_stream_function(X, Y, self.U, self.a)

        # 清除旧图形
        self.figure.clf()
        ax = self.figure.add_subplot(111, facecolor='black')  # 黑色背景坐标系

        # 绘制流线图（白色，带速度箭头）
        ax.streamplot(X, Y, Vx, Vy, 
                     color='white', linewidth=0.8,
                     density=1.5, arrowsize=0.8)

        # 绘制等势线（黄色虚线，15个层级）
        ax.contour(X, Y, psi, levels=15,
                  colors='yellow', linestyles='--',
                  linewidths=0.6, alpha=0.8)

        # 添加圆柱体表示（半透明红色圆形）
        ax.add_patch(plt.Circle((0, 0), self.a,
                              color='red', alpha=0.3))

        # 标注驻点位置（圆柱表面x轴两端）
        ax.plot([self.a, -self.a], [0, 0], 'ro', markersize=5)

        # 坐标轴设置
        ax.set(xlim=(-5, 5), ylim=(-5, 5),  # 固定显示范围
             xticks=[], yticks=[],          # 隐藏坐标刻度
             aspect='equal')               # 保持纵横比一致

        # 更新状态信息
        Re = 2 * self.U * self.a / nu  # 雷诺数计算公式（特征长度为2a）
        valid = validate_streamlines(self.a, self.U)
        status = (f"雷诺数 Re = {Re:.1e} | "
                 f"状态: {'✔ 验证通过' if valid else '✖ 验证失败'}")
        self.status_label.setText(status)
        # 根据验证结果设置颜色反馈
        self.status_label.setStyleSheet(
            "color: #0F0;" if valid else "color: #F00;")  

        self.canvas.draw()  # 触发画布重绘


# ================== 物理验证 ==================
def validate_streamlines(a, U):
    """
    流线闭合性验证（圆柱表面应为流线）
    
    参数:
        a (float): 圆柱半径，单位m
        U (float): 来流速度，单位m/s
    
    返回:
        bool: 是否通过验证（圆柱表面流函数标准差小于阈值）
    
    方法:
        1. 在圆柱表面均匀采样10个点
        2. 计算各点的流函数值
        3. 验证标准差是否小于阈值（U*a的1%）
    """
    theta = np.linspace(0, 2*np.pi, 10)  # 在圆周上均匀采样
    x = a * np.cos(theta)
    y = a * np.sin(theta)
    psi_values = calculate_stream_function(x, y, U, a)
    return np.std(psi_values) < 0.01 * U * a  # 允许1%的误差


# ================== 预编译和启动 ==================
if __name__ == '__main__':
    # Numba预编译（避免首次运行时的编译延迟）
    X_dummy = np.ones((2,2), dtype=np.float64)
    compute_velocity(X_dummy, X_dummy, 1.0, 0.5)  # 用虚拟数据触发编译
    
    # 启动Qt应用
    app = QApplication([])
    window = MainWindow()
    window.show()
    app.exec_()